{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XDYIwbXtHNvm"
   },
   "source": [
    "# Maps Tutorial\n",
    "\n",
    "This tutorial demonstrates the use of map data in the open dataset. It includes code to align laser point clouds to the map data and visualize the results.\n",
    "\n",
    "To run, use this [colab link](https://colab.research.google.com/github/waymo-research/waymo-open-dataset/blob/master/tutorial/tutorial_maps.ipynb) to open directly in colab."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "w5BBuR1JG3_Y"
   },
   "source": [
    "## Install Waymo Open Dataset Package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "t1vtxuAzGTTK"
   },
   "outputs": [],
   "source": [
    "!pip install waymo-open-dataset-tf-2-12-0==1.6.7"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zy79JF35iDaE"
   },
   "source": [
    "## Plot point clouds with maps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2MnJVU1jVe7a"
   },
   "outputs": [],
   "source": [
    "from typing import List\n",
    "\n",
    "import numpy as np\n",
    "import plotly.graph_objs as go\n",
    "from waymo_open_dataset import dataset_pb2\n",
    "from waymo_open_dataset.utils import frame_utils\n",
    "from waymo_open_dataset.utils import plot_maps\n",
    "\n",
    "def plot_point_clouds_with_maps(frames: List[dataset_pb2.Frame])->None:\n",
    "  \"\"\"Plot the point clouds within the given frames with map data.\n",
    "\n",
    "  Map data must be populated in the first frame in the list.\n",
    "\n",
    "  Args:\n",
    "    frames: A list of frames to be plotted, frames[0] must contain map data.\n",
    "  \"\"\"\n",
    "\n",
    "  # Plot the map features.\n",
    "  if len(frames) == 0:\n",
    "    return\n",
    "  figure = plot_maps.plot_map_features(frames[0].map_features)\n",
    "\n",
    "  for frame in frames:\n",
    "    # Parse the frame lidar data into range images.\n",
    "    range_images, camera_projections, seg_labels, range_image_top_poses = (\n",
    "        frame_utils.parse_range_image_and_camera_projection(frame)\n",
    "    )\n",
    "\n",
    "    # Project the range images into points.\n",
    "    points, cp_points = frame_utils.convert_range_image_to_point_cloud(\n",
    "        frame,\n",
    "        range_images,\n",
    "        camera_projections,\n",
    "        range_image_top_poses,\n",
    "        keep_polar_features=True,\n",
    "    )\n",
    "    xyz = points[0][:, 3:]\n",
    "    num_points = xyz.shape[0]\n",
    "    \n",
    "    # Transform the points from the vehicle frame to the world frame.\n",
    "    xyz = np.concatenate([xyz, np.ones([num_points, 1])], axis=-1)\n",
    "    transform = np.reshape(np.array(frame.pose.transform), [4, 4])\n",
    "    xyz = np.transpose(np.matmul(transform, np.transpose(xyz)))[:, 0:3]\n",
    "\n",
    "    # Correct the pose of the points into the coordinate system of the first\n",
    "    # frame to align with the map data.\n",
    "    offset = frame.map_pose_offset\n",
    "    points_offset = np.array([offset.x, offset.y, offset.z])\n",
    "    xyz += points_offset\n",
    "\n",
    "    # Plot the point cloud for this frame aligned with the map data.\n",
    "    intensity = points[0][:, 0]\n",
    "    figure.add_trace(\n",
    "        go.Scatter3d(\n",
    "            x=xyz[:, 0],\n",
    "            y=xyz[:, 1],\n",
    "            z=xyz[:, 2],\n",
    "            mode='markers',\n",
    "            marker=dict(\n",
    "                size=1,\n",
    "                color=intensity,  # set color to an array/list of desired values\n",
    "                colorscale='Pinkyl',  # choose a colorscale\n",
    "                opacity=0.8,\n",
    "            ),\n",
    "        )\n",
    "    )\n",
    "\n",
    "  figure.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5mhOt4geq12V"
   },
   "source": [
    "## Load frames data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "elBS3vh4qoqq"
   },
   "outputs": [],
   "source": [
    "!wget https://github.com/waymo-research/waymo-open-dataset/raw/master/tutorial/frames_with_maps.tfrecord"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WfRqjpyEt_Px"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "FILENAME = '/content/frames_with_maps.tfrecord'\n",
    "\n",
    "dataset = tf.data.TFRecordDataset(FILENAME, compression_type='')\n",
    "\n",
    "# Load only 2 frames. Note that using too many frames may be slow to display.\n",
    "frames = []\n",
    "count = 0\n",
    "for data in dataset:\n",
    "  frame = dataset_pb2.Frame.FromString(bytearray(data.numpy()))\n",
    "  frames.append(frame)\n",
    "  count += 1\n",
    "  if count == 2: \n",
    "    break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RzTEio68tE2B"
   },
   "source": [
    "## Plot frames data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-Y59EuqTs8c1"
   },
   "outputs": [],
   "source": [
    "# Interactive plot of multiple point clouds aligned to the maps frame.\n",
    "\n",
    "# For most systems:\n",
    "#   left mouse button:   rotate\n",
    "#   right mouse button:  pan\n",
    "#   scroll wheel:        zoom\n",
    "\n",
    "plot_point_clouds_with_maps(frames)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "private_outputs": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
